

<!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  
  <title>mindspore.ops.operations.inner_ops &mdash; MindSpore master documentation</title>
  

  
  <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />

  
  

  
  

  

  
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/language_data.js"></script>
        <script async="async" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>
    
    <script type="text/javascript" src="../../../../_static/js/theme.js"></script>

    
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../../../index.html" class="icon icon-home"> MindSpore
          

          
          </a>

          
            
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption"><span class="caption-text">MindSpore Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.html">mindspore</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.common.initializer.html">mindspore.common.initializer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.communication.html">mindspore.communication</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.compression.html">mindspore.compression</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.context.html">mindspore.context</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.html">mindspore.dataset</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.audio.html">mindspore.dataset.audio</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.config.html">mindspore.dataset.config</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.text.html">mindspore.dataset.text</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.transforms.html">mindspore.dataset.transforms</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.vision.html">mindspore.dataset.vision</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.mindrecord.html">mindspore.mindrecord</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.html">mindspore.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.probability.html">mindspore.nn.probability</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.transformer.html">mindspore.nn.transformer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.numpy.html">mindspore.numpy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.ops.html">mindspore.ops</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.html">mindspore.parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.nn.html">mindspore.parallel.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.profiler.html">mindspore.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.scipy.html">mindspore.scipy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.train.html">mindspore.train</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.boost.html">mindspore.boost</a></li>
</ul>
<p class="caption"><span class="caption-text">MindSpore C++ API</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://www.mindspore.cn/lite/api/zh-CN/master/api_cpp/mindspore.html">MindSpore Lite↗</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">

      
      <nav class="wy-nav-top" aria-label="top navigation">
        
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">MindSpore</a>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        
          

















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
        
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
        
      <li>mindspore.ops.operations.inner_ops</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <h1>Source code for mindspore.ops.operations.inner_ops</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2020 Huawei Technologies Co., Ltd</span>
<span class="c1">#</span>
<span class="c1"># Licensed under the Apache License, Version 2.0 (the &quot;License&quot;);</span>
<span class="c1"># you may not use this file except in compliance with the License.</span>
<span class="c1"># You may obtain a copy of the License at</span>
<span class="c1">#</span>
<span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span>
<span class="c1">#</span>
<span class="c1"># Unless required by applicable law or agreed to in writing, software</span>
<span class="c1"># distributed under the License is distributed on an &quot;AS IS&quot; BASIS,</span>
<span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span>
<span class="c1"># See the License for the specific language governing permissions and</span>
<span class="c1"># limitations under the License.</span>
<span class="c1"># ============================================================================</span>

<span class="sd">&quot;&quot;&quot;inner_ops&quot;&quot;&quot;</span>

<span class="kn">import</span> <span class="nn">numbers</span>
<span class="kn">from</span> <span class="nn">..._checkparam</span> <span class="kn">import</span> <span class="n">Validator</span> <span class="k">as</span> <span class="n">validator</span>
<span class="kn">from</span> <span class="nn">..._checkparam</span> <span class="kn">import</span> <span class="n">Rel</span>
<span class="kn">from</span> <span class="nn">...common</span> <span class="kn">import</span> <span class="n">dtype</span> <span class="k">as</span> <span class="n">mstype</span>
<span class="kn">from</span> <span class="nn">...common.dtype</span> <span class="kn">import</span> <span class="n">tensor</span><span class="p">,</span> <span class="n">dtype_to_pytype</span>
<span class="kn">from</span> <span class="nn">..primitive</span> <span class="kn">import</span> <span class="n">prim_attr_register</span><span class="p">,</span> <span class="n">Primitive</span><span class="p">,</span> <span class="n">PrimitiveWithInfer</span>
<span class="kn">from</span> <span class="nn">..</span> <span class="kn">import</span> <span class="n">signature</span> <span class="k">as</span> <span class="n">sig</span>


<div class="viewcode-block" id="ScalarCast"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.ScalarCast.html#mindspore.ops.ScalarCast">[docs]</a><span class="k">class</span> <span class="nc">ScalarCast</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Casts the input scalar to another type.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **input_x** (scalar) - The input scalar. Only constant value is allowed.</span>
<span class="sd">        - **input_y** (mindspore.dtype) - The type to be cast. Only constant value is allowed.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Scalar. The type is the same as the python type corresponding to `input_y`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If neither `input_x` nor `input_y` is a constant value.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; scalar_cast = ops.ScalarCast()</span>
<span class="sd">        &gt;&gt;&gt; output = scalar_cast(255.0, mindspore.int32)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        255</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">pass</span>

    <span class="k">def</span> <span class="nf">__infer__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">t</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_equal_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">x</span><span class="p">[</span><span class="s1">&#39;shape&#39;</span><span class="p">]),</span> <span class="mi">0</span><span class="p">,</span> <span class="s1">&#39;x shape&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">value</span><span class="p">,</span> <span class="n">to</span> <span class="o">=</span> <span class="n">x</span><span class="p">[</span><span class="s1">&#39;value&#39;</span><span class="p">],</span> <span class="n">t</span><span class="p">[</span><span class="s1">&#39;value&#39;</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">value</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;value&quot;</span><span class="p">,</span> <span class="n">value</span><span class="p">,</span> <span class="p">[</span><span class="n">numbers</span><span class="o">.</span><span class="n">Number</span><span class="p">,</span> <span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">to</span><span class="p">,</span> <span class="nb">type</span><span class="p">(</span><span class="n">tensor</span><span class="p">)):</span>
                <span class="n">to</span> <span class="o">=</span> <span class="n">to</span><span class="o">.</span><span class="n">element_type</span><span class="p">()</span>
            <span class="n">np_type</span> <span class="o">=</span> <span class="n">dtype_to_pytype</span><span class="p">(</span><span class="n">to</span><span class="p">)</span>
            <span class="n">value</span> <span class="o">=</span> <span class="n">np_type</span><span class="p">(</span><span class="n">value</span><span class="p">)</span>
        <span class="n">out</span> <span class="o">=</span> <span class="p">{</span><span class="s1">&#39;shape&#39;</span><span class="p">:</span> <span class="n">x</span><span class="p">[</span><span class="s1">&#39;shape&#39;</span><span class="p">],</span>
               <span class="s1">&#39;dtype&#39;</span><span class="p">:</span> <span class="n">t</span><span class="p">[</span><span class="s1">&#39;value&#39;</span><span class="p">],</span>
               <span class="s1">&#39;value&#39;</span><span class="p">:</span> <span class="n">value</span><span class="p">}</span>
        <span class="k">return</span> <span class="n">out</span></div>


<div class="viewcode-block" id="Randperm"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.Randperm.html#mindspore.ops.Randperm">[docs]</a><span class="k">class</span> <span class="nc">Randperm</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Generates n random samples from 0 to n-1 without repeating. If `max_length` &gt; n,</span>
<span class="sd">    the last `max_length-n` elements will be filled with `pad`.</span>

<span class="sd">    Args:</span>
<span class="sd">        max_length (int): Number of items expected to get and the number must be greater than 0. Default: 1.</span>
<span class="sd">        pad (int): The pad value to be filled. Default: -1.</span>
<span class="sd">        dtype (mindspore.dtype): The type of output. Default: mindspore.int32.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **n** (Tensor[int32]) - The input tensor with shape: (1,) and the number must be in [0, `max_length`].</span>

<span class="sd">    Outputs:</span>
<span class="sd">        - **output** (Tensor) - The output Tensor with shape: (`max_length`,) and type: `dtype`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If neither `max_length` nor `pad` is an int.</span>
<span class="sd">        TypeError: If `n` is not a Tensor.</span>
<span class="sd">        TypeError: If `n` has non-Int elements.</span>
<span class="sd">        TypeError: If `n` has negative elements.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; # The result of every execution is different because this operator will generate n random samples.</span>
<span class="sd">        &gt;&gt;&gt; randperm = ops.Randperm(max_length=30, pad=-1)</span>
<span class="sd">        &gt;&gt;&gt; n = Tensor([20], dtype=mindspore.int32)</span>
<span class="sd">        &gt;&gt;&gt; output = randperm(n)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [15 6 11 19 14 16 9 5 13 18 4 10 8 0 17 2 1 12 3 7</span>
<span class="sd">         -1 -1 -1 -1 -1 -1 -1 -1 -1 -1]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">max_length</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">pad</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">mstype</span><span class="o">.</span><span class="n">int32</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize Randperm&quot;&quot;&quot;</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;pad&quot;</span><span class="p">,</span> <span class="n">pad</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;max_length&quot;</span><span class="p">,</span> <span class="n">max_length</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_int</span><span class="p">(</span><span class="n">max_length</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">GE</span><span class="p">,</span> <span class="s2">&quot;max_length&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">dtype</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_length</span> <span class="o">=</span> <span class="n">max_length</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">init_prim_io_names</span><span class="p">(</span><span class="n">inputs</span><span class="o">=</span><span class="p">[],</span> <span class="n">outputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;output&#39;</span><span class="p">])</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">n_shape</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">n_shape</span><span class="p">),</span> <span class="mi">1</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="s2">&quot;rank_of_n&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_int</span><span class="p">(</span><span class="n">n_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="mi">1</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="s2">&quot;length_of_n&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">max_length</span><span class="p">]</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">n_type</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_type_name</span><span class="p">(</span><span class="s2">&quot;n_type&quot;</span><span class="p">,</span> <span class="n">n_type</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">int32</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>

        <span class="n">valid_values</span> <span class="o">=</span> <span class="p">(</span><span class="n">mstype</span><span class="o">.</span><span class="n">int8</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">int16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">int32</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">int64</span><span class="p">,</span>
                        <span class="n">mstype</span><span class="o">.</span><span class="n">uint8</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">uint16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">uint32</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">uint64</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_type_name</span><span class="p">(</span><span class="s2">&quot;dtype&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">valid_values</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span></div>


<div class="viewcode-block" id="NoRepeatNGram"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.NoRepeatNGram.html#mindspore.ops.NoRepeatNGram">[docs]</a><span class="k">class</span> <span class="nc">NoRepeatNGram</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Updates log_probs with repeat n-grams.</span>

<span class="sd">    During beam search, if consecutive `ngram_size` words exist in the generated word sequence,</span>
<span class="sd">    the consecutive `ngram_size` words will be avoided during subsequent prediction.</span>
<span class="sd">    For example, when `ngram_size` is 3, the generated word sequence is [1, 2, 3, 2, 3],</span>
<span class="sd">    the next predicted word will not be 2 and the value of `log_probs` will be replaced with -FLOAT_MAX.</span>
<span class="sd">    Because 3 consecutive words [2, 3, 2] do not appear twice in the word sequence.</span>

<span class="sd">    Args:</span>
<span class="sd">        ngram_size (int): Size of n-grams, must be greater than 0. Default: 1.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **state_seq** (Tensor) - A 3-D tensor with shape: (batch_size, beam_width, m).</span>
<span class="sd">        - **log_probs** (Tensor) - A 3-D tensor with shape: (batch_size, beam_width, vocab_size).</span>
<span class="sd">          The value of log_probs will be replaced with -FLOAT_MAX when n-grams repeated.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        - **log_probs** (Tensor) - The output Tensor with same shape and type as original `log_probs`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `ngram_size` is not an int.</span>
<span class="sd">        TypeError: If neither `state_seq` nor `log_probs` is a Tensor.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; no_repeat_ngram = ops.NoRepeatNGram(ngram_size=3)</span>
<span class="sd">        &gt;&gt;&gt; state_seq = Tensor([[[1, 2, 1, 2, 5, 1, 2],</span>
<span class="sd">        ...                      [9, 3, 9, 5, 4, 1, 5]],</span>
<span class="sd">        ...                     [[4, 8, 6, 4, 5, 6, 4],</span>
<span class="sd">        ...                      [4, 8, 8, 4, 3, 4, 8]]], dtype=mindspore.int32)</span>
<span class="sd">        &gt;&gt;&gt; log_probs = Tensor([[[0.7, 0.8, 0.6, 0.9, 0.2, 0.8, 0.4, 0.6, 0.2, 0.7],</span>
<span class="sd">        ...                      [0.4, 0.5, 0.6, 0.7, 0.8, 0.1, 0.9, 0.8, 0.7, 0.1]],</span>
<span class="sd">        ...                     [[0.9, 0.7, 0.6, 0.3, 0.5, 0.3, 0.5, 0.4, 0.8, 0.6],</span>
<span class="sd">        ...                      [0.5, 0.8, 0.8, 0.7, 0.7, 0.8, 0.2, 0.7, 0.9, 0.7]]], dtype=mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; output = no_repeat_ngram(state_seq, log_probs)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[[ 6.9999999e-01 -3.4028235e+38  6.0000002e-01  8.9999998e-01</span>
<span class="sd">            2.0000000e-01 -3.4028235e+38  4.0000001e-01  6.0000002e-01</span>
<span class="sd">            2.0000000e-01  6.9999999e-01]</span>
<span class="sd">          [ 4.0000001e-01  5.0000000e-01  6.0000002e-01  6.9999999e-01</span>
<span class="sd">            8.0000001e-01  1.0000000e-01  8.9999998e-01  8.0000001e-01</span>
<span class="sd">            6.9999999e-01  1.0000000e-01]]</span>
<span class="sd">         [[ 8.9999998e-01  6.9999999e-01  6.0000002e-01  3.0000001e-01</span>
<span class="sd">            5.0000000e-01 -3.4028235e+38  5.0000000e-01  4.0000001e-01</span>
<span class="sd">            8.0000001e-01  6.0000002e-01]</span>
<span class="sd">          [ 5.0000000e-01  8.0000001e-01  8.0000001e-01  6.9999999e-01</span>
<span class="sd">            6.9999999e-01  8.0000001e-01  2.0000000e-01  6.9999999e-01</span>
<span class="sd">           -3.4028235e+38  6.9999999e-01]]]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ngram_size</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;NoRepeatNGram Randperm&quot;&quot;&quot;</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;ngram_size&quot;</span><span class="p">,</span> <span class="n">ngram_size</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_int</span><span class="p">(</span><span class="n">ngram_size</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">GE</span><span class="p">,</span> <span class="s2">&quot;ngram_size&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ngram_size</span> <span class="o">=</span> <span class="n">ngram_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">init_prim_io_names</span><span class="p">(</span><span class="n">inputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;state_seq&#39;</span><span class="p">,</span> <span class="s1">&#39;log_probs&#39;</span><span class="p">],</span> <span class="n">outputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;log_probs&#39;</span><span class="p">])</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">seq_shape</span><span class="p">,</span> <span class="n">log_shape</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">seq_shape</span><span class="p">),</span> <span class="mi">3</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="s2">&quot;rank of state_seq&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">log_shape</span><span class="p">),</span> <span class="mi">3</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="s2">&quot;rank of log_probs&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;state_seq shape[0]&quot;</span><span class="p">,</span> <span class="n">seq_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="s2">&quot;log_probs shape[0]&quot;</span><span class="p">,</span> <span class="n">log_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;state_seq shape[1]&quot;</span><span class="p">,</span> <span class="n">seq_shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="s2">&quot;log_probs shape[1]&quot;</span><span class="p">,</span> <span class="n">log_shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;ngram_size&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">ngram_size</span><span class="p">,</span> <span class="s2">&quot;state_seq shape[2] + 1&quot;</span><span class="p">,</span> <span class="n">seq_shape</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">LE</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">log_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">seq_type</span><span class="p">,</span> <span class="n">log_type</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_type_name</span><span class="p">(</span><span class="s2">&quot;seq_type&quot;</span><span class="p">,</span> <span class="n">seq_type</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">int32</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">valid_values</span> <span class="o">=</span> <span class="p">(</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float64</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_type_name</span><span class="p">(</span><span class="s2">&quot;log_type&quot;</span><span class="p">,</span> <span class="n">log_type</span><span class="p">,</span> <span class="n">valid_values</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">log_type</span></div>


<span class="k">class</span> <span class="nc">LambApplyOptimizerAssign</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Updates gradients by LAMB optimizer algorithm. Get the compute ratio.</span>

<span class="sd">    The Lamb optimizer is proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes</span>
<span class="sd">    &lt;https://arxiv.org/abs/1904.00962&gt;`_.</span>

<span class="sd">    The updating formulas are as follows,</span>

<span class="sd">    .. math::</span>
<span class="sd">        \begin{array}{ll} \\</span>
<span class="sd">            m = \beta_1 * m + (1 - \beta_1) * g \\</span>
<span class="sd">            v = \beta_2 * v + (1 - \beta_2) * g * g \\</span>
<span class="sd">            m = \frac{m}{1 - \beta_1^t} \\</span>
<span class="sd">            v = \frac{v}{1 - \beta_2^t} \\</span>
<span class="sd">            r = \frac{m}{\sqrt{v} + \epsilon} \\</span>
<span class="sd">            w = w - l * \frac{\left \| w \right \|}{\left \| r \right \|} * (r + \lambda * w))</span>
<span class="sd">        \end{array}</span>

<span class="sd">    :math:`m` represents the 1st moment vector, :math:`v` represents the 2nd moment vector, :math:`g` represents</span>
<span class="sd">    `gradient`, :math:`l` represents learning rate `lr`, :math:`\beta_1, \beta_2` represent `beta1` and `beta2`,</span>
<span class="sd">    :math:`t` represents updating step while :math:`beta_1^t` and :math:`beta_2^t` represent `beta1_power` and</span>
<span class="sd">    `beta2_power`, :math:`\lambda` represents `weight_decay`, :math:`w` represents `var`, :math:`\epsilon` represents</span>
<span class="sd">    `epsilon`.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **gradient** (Tensor) - Gradient of parameters, float32/float16.</span>
<span class="sd">        - **v** (Tensor) - the 2nd moment vector in the updating formula, has the same type as `gradient`.</span>
<span class="sd">        - **m** (Tensor) - The 1st moment vector in the updating formula, has the same type as `gradient`.</span>
<span class="sd">        - **var** (Tensor) - Weights to be updated, has the same type as `gradient`.</span>
<span class="sd">        - **beta1** (Tensor) - :math:`beta_1` in the updating formula, float32/float16.</span>
<span class="sd">        - **sub1** (Tensor) - :math:`1-beta_1` in the updating formula, has the same type as `beta1`.</span>
<span class="sd">        - **beta2** (Tensor) - :math:`beta_2` in the updating formula, has the same type as `beta1`.</span>
<span class="sd">        - **sub2** (Tensor) - :math:`1-beta_2` in the updating formula, has the same type as `beta1`.</span>
<span class="sd">        - **epsilon** (Tensor) - Term added to the denominator, has the same type as `beta1`.</span>
<span class="sd">        - **steps** (Tensor) - :math:`t` in the updating formula, global step, has the same type as `beta1`.</span>
<span class="sd">        - **lr** (Tensor) - :math:`l` in the updating formula, learning rate, has the same type as `beta1`.</span>
<span class="sd">        - **decay_flag** (Tensor) -Specify whether param update with weight decay, has the same type as `beta1`.</span>
<span class="sd">        - **weight_decay** (Tensor) - :math:`\lambda` in the updating formula, has the same type as `beta1`.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, the compute ratio r.</span>
<span class="sd">        - **update** (Tensor) - :math:`r + \lambda * w` in the updating formula. The same shape and data type as `m`.</span>
<span class="sd">        - **v** (Tensor) - the 2nd moment vector in the updating formula after updated inplace,</span>
<span class="sd">                           has the same type as `gradient`.</span>
<span class="sd">        - **m** (Tensor) - The 1st moment vector in the updating formula after updated inplace,</span>
<span class="sd">                           has the same type as `gradient`.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend``</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize LambApplyOptimizerAssign&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;side_effect_mem&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">grad_shape</span><span class="p">,</span> <span class="n">v_shape</span><span class="p">,</span> <span class="n">m_shape</span><span class="p">,</span> <span class="n">var_shape</span><span class="p">,</span> <span class="n">beta1_shape</span><span class="p">,</span> <span class="n">sub1_shape</span><span class="p">,</span>
                    <span class="n">beta2_shape</span><span class="p">,</span> <span class="n">sub2_shape</span><span class="p">,</span> <span class="n">eps_shape</span><span class="p">,</span> <span class="n">steps_shape</span><span class="p">,</span> <span class="n">use_weight_shape</span><span class="p">,</span> <span class="n">weight_decay_shape</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;var_shape&quot;</span><span class="p">,</span> <span class="n">var_shape</span><span class="p">,</span> <span class="s2">&quot;m_shape&quot;</span><span class="p">,</span> <span class="n">m_shape</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;var_shape&quot;</span><span class="p">,</span> <span class="n">var_shape</span><span class="p">,</span> <span class="s2">&quot;v_shape&quot;</span><span class="p">,</span> <span class="n">v_shape</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;var_shape&quot;</span><span class="p">,</span> <span class="n">var_shape</span><span class="p">,</span> <span class="s2">&quot;grad_shape&quot;</span><span class="p">,</span> <span class="n">grad_shape</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">m_shape</span><span class="p">,</span> <span class="n">v_shape</span><span class="p">,</span> <span class="n">m_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">grad_dtype</span><span class="p">,</span> <span class="n">v_dtype</span><span class="p">,</span> <span class="n">m_dtype</span><span class="p">,</span> <span class="n">var_dtype</span><span class="p">,</span> <span class="n">beta1_dtype</span><span class="p">,</span> <span class="n">sub1_dtype</span><span class="p">,</span>
                    <span class="n">beta2_dtype</span><span class="p">,</span> <span class="n">sub2_dtype</span><span class="p">,</span> <span class="n">eps_dtype</span><span class="p">,</span> <span class="n">steps_dtype</span><span class="p">,</span> <span class="n">use_weight_dtype</span><span class="p">,</span> <span class="n">weight_decay_dtype</span><span class="p">):</span>
        <span class="n">args</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;var&quot;</span><span class="p">:</span> <span class="n">var_dtype</span><span class="p">,</span> <span class="s2">&quot;m&quot;</span><span class="p">:</span> <span class="n">m_dtype</span><span class="p">,</span> <span class="s2">&quot;v&quot;</span><span class="p">:</span> <span class="n">v_dtype</span><span class="p">,</span> <span class="s2">&quot;grad&quot;</span><span class="p">:</span> <span class="n">grad_dtype</span><span class="p">}</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensors_dtypes_same_and_valid</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>

        <span class="n">args</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;beta1&quot;</span><span class="p">:</span> <span class="n">beta1_dtype</span><span class="p">,</span> <span class="s2">&quot;sub1&quot;</span><span class="p">:</span> <span class="n">sub1_dtype</span><span class="p">,</span> <span class="s2">&quot;beta2&quot;</span><span class="p">:</span> <span class="n">beta2_dtype</span><span class="p">,</span> <span class="s2">&quot;sub2&quot;</span><span class="p">:</span> <span class="n">sub2_dtype</span><span class="p">,</span>
                <span class="s2">&quot;eps&quot;</span><span class="p">:</span> <span class="n">eps_dtype</span><span class="p">,</span> <span class="s2">&quot;steps&quot;</span><span class="p">:</span> <span class="n">steps_dtype</span><span class="p">,</span> <span class="s2">&quot;use_weight&quot;</span><span class="p">:</span> <span class="n">use_weight_dtype</span><span class="p">,</span>
                <span class="s2">&quot;weight_decay&quot;</span><span class="p">:</span> <span class="n">weight_decay_dtype</span><span class="p">}</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_scalar_or_tensor_types_same</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">m_dtype</span><span class="p">,</span> <span class="n">v_dtype</span><span class="p">,</span> <span class="n">v_dtype</span>


<span class="k">class</span> <span class="nc">LambApplyWeightAssign</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Updates gradients by LAMB optimizer algorithm. The weight update part.</span>

<span class="sd">    The Lamb optimizer is proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes</span>
<span class="sd">    &lt;https://arxiv.org/abs/1904.00962&gt;`_.</span>

<span class="sd">    The updating formulas are as follows,</span>

<span class="sd">    .. math::</span>
<span class="sd">        \begin{array}{ll} \\</span>
<span class="sd">            m = \beta_1 * m + (1 - \beta_1) * g \\</span>
<span class="sd">            v = \beta_2 * v + (1 - \beta_2) * g * g \\</span>
<span class="sd">            m = \frac{m}{1 - \beta_1^t} \\</span>
<span class="sd">            v = \frac{v}{1 - \beta_2^t} \\</span>
<span class="sd">            r = \frac{m}{\sqrt{v} + \epsilon} \\</span>
<span class="sd">            w = w - l * \frac{\left \| w \right \|}{\left \| r \right \|} * (r + \lambda * w))</span>
<span class="sd">        \end{array}</span>

<span class="sd">    :math:`m` represents the 1st moment vector, :math:`v` represents the 2nd moment vector, :math:`g` represents</span>
<span class="sd">    `gradient`, :math:`l` represents learning rate `lr`, :math:`\beta_1, \beta_2` represent `beta1` and `beta2`,</span>
<span class="sd">    :math:`t` represents updating step while :math:`beta_1^t` and :math:`beta_2^t` represent `beta1_power` and</span>
<span class="sd">    `beta2_power`, :math:`\lambda` represents `weight_decay`, :math:`w` represents `var`, :math:`\epsilon` represents</span>
<span class="sd">    `epsilon`.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **w_norm** (Tensor) - :math:`\left \| w \right \|` in the updating formula, float32/float16.</span>
<span class="sd">        - **g_norm** (Tensor) - :math:`\left \| r \right \|` in the updating formula, has the same type as `w_norm`.</span>
<span class="sd">        - **lr** (Tensor) - :math:`l` in the updating formula, the learning rate, float32/float16.</span>
<span class="sd">        - **update** (Tensor) -:math:`r + \lambda * w`in the updating formula, float32/float16.</span>
<span class="sd">        - **var** (Tensor) - Weights to be updated, the same shape and type as `update`.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        - **var** (Tensor) - Weights to be updated in place, the same shape and type as `var` in inputs.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend``</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize LambApplyWeightAssign&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;side_effect_mem&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">w_norm_shape</span><span class="p">,</span> <span class="n">g_norm_shape</span><span class="p">,</span> <span class="n">lr_shape</span><span class="p">,</span> <span class="n">update_shape</span><span class="p">,</span> <span class="n">var_shape</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;var_shape&quot;</span><span class="p">,</span> <span class="n">var_shape</span><span class="p">,</span> <span class="s2">&quot;update_shape&quot;</span><span class="p">,</span> <span class="n">update_shape</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">var_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">w_norm_dtype</span><span class="p">,</span> <span class="n">g_norm_dtype</span><span class="p">,</span> <span class="n">lr_dtype</span><span class="p">,</span> <span class="n">update_dtype</span><span class="p">,</span> <span class="n">var_dtype</span><span class="p">):</span>
        <span class="n">args</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;var&quot;</span><span class="p">:</span> <span class="n">var_dtype</span><span class="p">,</span> <span class="s2">&quot;update&quot;</span><span class="p">:</span> <span class="n">update_dtype</span><span class="p">}</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensors_dtypes_same_and_valid</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>

        <span class="n">args</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;w_norm&quot;</span><span class="p">:</span> <span class="n">w_norm_dtype</span><span class="p">,</span> <span class="s2">&quot;g_norm&quot;</span><span class="p">:</span> <span class="n">g_norm_dtype</span><span class="p">,</span> <span class="s2">&quot;lr&quot;</span><span class="p">:</span> <span class="n">lr_dtype</span><span class="p">}</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_scalar_or_tensor_types_same</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">var_dtype</span>


<span class="k">class</span> <span class="nc">MakeRefKey</span><span class="p">(</span><span class="n">Primitive</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Makes a RefKey instance by string. RefKey stores the name of Parameter, can be passed through the functions,</span>
<span class="sd">    and used for Assign target.</span>

<span class="sd">    Args:</span>
<span class="sd">        tag (str): Parameter name to make the RefKey.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        No inputs.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        RefKeyType, made from the Parameter name.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `tag` is not a str.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Parameter, Tensor</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import dtype as mstype</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.y = Parameter(Tensor(np.ones([2, 3]), mstype.int32), name=&quot;y&quot;)</span>
<span class="sd">        ...         self.make_ref_key = ops.MakeRefKey(&quot;y&quot;)</span>
<span class="sd">        ...</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         key = self.make_ref_key()</span>
<span class="sd">        ...         ref = ops.make_ref(key, x, self.y)</span>
<span class="sd">        ...         return ref * x</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.array([[1, 2, 3], [4, 5, 6]]), mindspore.int32)</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; output = net(x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[ 1  4  9]</span>
<span class="sd">         [16 25 36]]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tag</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;tag&#39;</span><span class="p">,</span> <span class="n">tag</span><span class="p">,</span> <span class="p">(</span><span class="nb">str</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">pass</span>


<span class="k">class</span> <span class="nc">FusedWeightScaleApplyMomentum</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Optimizer that implements the Momentum algorithm with weight decay and loss scale.</span>

<span class="sd">    Refer to the paper `On the importance of initialization and momentum in deep</span>
<span class="sd">    learning &lt;https://dl.acm.org/doi/10.5555/3042817.3043064&gt;`_  for more details.</span>

<span class="sd">    Refer to :class:`mindspore.nn.Momentum` for more details about the formula and usage.</span>

<span class="sd">    Inputs of `variable`, `accumulation` and `gradient` comply with the implicit type conversion rules</span>
<span class="sd">    to make the data types consistent.</span>
<span class="sd">    If they have different data types, the lower priority data type will be converted to</span>
<span class="sd">    relatively highest priority data type.</span>
<span class="sd">    Data type conversion of Parameter is not supported. RuntimeError exception will be thrown.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **weight_decay** (Tensor) - The weight decay value, must be a scalar tensor with float data type.</span>
<span class="sd">          Default: 0.0.</span>
<span class="sd">        - **loss_scale** (Tensor) - The loss scale value, must be a scalar tensor with float data type.</span>
<span class="sd">          Default: 1.0.</span>
<span class="sd">        - **variable** (Parameter) - Weights to be updated. data type must be float.</span>
<span class="sd">        - **accumulation** (Parameter) - Accumulated gradient value by moment weight.</span>
<span class="sd">          Has the same data type with `variable`.</span>
<span class="sd">        - **learning_rate** (Union[Number, Tensor]) - The learning rate value, must be a float number or</span>
<span class="sd">          a scalar tensor with float data type.</span>
<span class="sd">        - **gradient** (Tensor) - Gradient, has the same data type as `variable`.</span>
<span class="sd">        - **momentum** (Union[Number, Tensor]) - Momentum, must be a float number or</span>
<span class="sd">          a scalar tensor with float data type.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, parameters to be updated.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``GPU``</span>
<span class="sd">    Examples:</span>
<span class="sd">        Please refer to the usage in :class:`mindspore.nn.Momentum`, and add weight_decay and loss_scale as inputs.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">__mindspore_signature__</span> <span class="o">=</span> <span class="p">(</span>
        <span class="n">sig</span><span class="o">.</span><span class="n">make_sig</span><span class="p">(</span><span class="s1">&#39;weight_decay&#39;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">sig</span><span class="o">.</span><span class="n">sig_dtype</span><span class="o">.</span><span class="n">T3</span><span class="p">),</span>
        <span class="n">sig</span><span class="o">.</span><span class="n">make_sig</span><span class="p">(</span><span class="s1">&#39;loss_scale&#39;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">sig</span><span class="o">.</span><span class="n">sig_dtype</span><span class="o">.</span><span class="n">T3</span><span class="p">),</span>
        <span class="n">sig</span><span class="o">.</span><span class="n">make_sig</span><span class="p">(</span><span class="s1">&#39;variable&#39;</span><span class="p">,</span> <span class="n">sig</span><span class="o">.</span><span class="n">sig_rw</span><span class="o">.</span><span class="n">RW_WRITE</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">sig</span><span class="o">.</span><span class="n">sig_dtype</span><span class="o">.</span><span class="n">T</span><span class="p">),</span>
        <span class="n">sig</span><span class="o">.</span><span class="n">make_sig</span><span class="p">(</span><span class="s1">&#39;accumulation&#39;</span><span class="p">,</span> <span class="n">sig</span><span class="o">.</span><span class="n">sig_rw</span><span class="o">.</span><span class="n">RW_WRITE</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">sig</span><span class="o">.</span><span class="n">sig_dtype</span><span class="o">.</span><span class="n">T</span><span class="p">),</span>
        <span class="n">sig</span><span class="o">.</span><span class="n">make_sig</span><span class="p">(</span><span class="s1">&#39;learning_rate&#39;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">sig</span><span class="o">.</span><span class="n">sig_dtype</span><span class="o">.</span><span class="n">T1</span><span class="p">),</span>
        <span class="n">sig</span><span class="o">.</span><span class="n">make_sig</span><span class="p">(</span><span class="s1">&#39;gradient&#39;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">sig</span><span class="o">.</span><span class="n">sig_dtype</span><span class="o">.</span><span class="n">T</span><span class="p">),</span>
        <span class="n">sig</span><span class="o">.</span><span class="n">make_sig</span><span class="p">(</span><span class="s1">&#39;momentum&#39;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">sig</span><span class="o">.</span><span class="n">sig_dtype</span><span class="o">.</span><span class="n">T2</span><span class="p">)</span>
    <span class="p">)</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">init_prim_io_names</span><span class="p">(</span><span class="n">inputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;weight_decay&#39;</span><span class="p">,</span> <span class="s1">&#39;loss_scale&#39;</span><span class="p">,</span> <span class="s1">&#39;variable&#39;</span><span class="p">,</span> <span class="s1">&#39;accumulation&#39;</span><span class="p">,</span> <span class="s1">&#39;learning_rate&#39;</span><span class="p">,</span>
                                        <span class="s1">&#39;gradient&#39;</span><span class="p">,</span> <span class="s1">&#39;momentum&#39;</span><span class="p">],</span> <span class="n">outputs</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;output&#39;</span><span class="p">])</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">d_shape</span><span class="p">,</span> <span class="n">s_shape</span><span class="p">,</span> <span class="n">v_shape</span><span class="p">,</span> <span class="n">a_shape</span><span class="p">,</span> <span class="n">l_shape</span><span class="p">,</span> <span class="n">g_shape</span><span class="p">,</span> <span class="n">m_shape</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">v_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">d_dtype</span><span class="p">,</span> <span class="n">s_dtype</span><span class="p">,</span> <span class="n">v_dtype</span><span class="p">,</span> <span class="n">a_dtype</span><span class="p">,</span> <span class="n">l_dtype</span><span class="p">,</span> <span class="n">g_dtype</span><span class="p">,</span> <span class="n">m_dtype</span><span class="p">):</span>
        <span class="n">valid_dtypes</span> <span class="o">=</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">v_dtype</span> <span class="o">!=</span> <span class="n">mstype</span><span class="o">.</span><span class="n">type_refkey</span> <span class="ow">and</span> <span class="n">a_dtype</span> <span class="o">!=</span> <span class="n">mstype</span><span class="o">.</span><span class="n">type_refkey</span><span class="p">:</span>
            <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s2">&quot;v&quot;</span><span class="p">,</span> <span class="n">v_dtype</span><span class="p">,</span> <span class="n">valid_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
            <span class="n">validator</span><span class="o">.</span><span class="n">check_tensor_dtype_valid</span><span class="p">(</span><span class="s2">&quot;a&quot;</span><span class="p">,</span> <span class="n">a_dtype</span><span class="p">,</span> <span class="n">valid_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_scalar_or_tensor_types_same</span><span class="p">({</span><span class="s2">&quot;l_dtype&quot;</span><span class="p">:</span> <span class="n">l_dtype</span><span class="p">},</span> <span class="n">valid_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_scalar_or_tensor_types_same</span><span class="p">({</span><span class="s2">&quot;g_dtype&quot;</span><span class="p">:</span> <span class="n">g_dtype</span><span class="p">},</span> <span class="n">valid_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_scalar_or_tensor_types_same</span><span class="p">({</span><span class="s2">&quot;m_dtype&quot;</span><span class="p">:</span> <span class="n">m_dtype</span><span class="p">},</span> <span class="n">valid_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_scalar_or_tensor_types_same</span><span class="p">({</span><span class="s2">&quot;d_dtype&quot;</span><span class="p">:</span> <span class="n">d_dtype</span><span class="p">},</span> <span class="n">valid_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_scalar_or_tensor_types_same</span><span class="p">({</span><span class="s2">&quot;s_dtype&quot;</span><span class="p">:</span> <span class="n">s_dtype</span><span class="p">},</span> <span class="n">valid_dtypes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">v_dtype</span>


<span class="k">class</span> <span class="nc">FusedCastAdamWeightDecay</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Updates gradients by the Adaptive Moment Estimation (AdamWeightDecay) algorithm with weight decay. This operator</span>
<span class="sd">    incorporates type conversion when parameters are initialized with dtype of float16.</span>

<span class="sd">    The Adam algorithm is proposed in `Adam: A Method for Stochastic Optimization &lt;https://arxiv.org/abs/1412.6980&gt;`_.</span>
<span class="sd">    The AdamWeightDecay variant was proposed in `Decoupled Weight Decay Regularization</span>
<span class="sd">    &lt;https://arxiv.org/abs/1711.05101&gt;`_.</span>

<span class="sd">    The updating formulas are as follows,</span>

<span class="sd">    .. math::</span>
<span class="sd">        \begin{array}{ll} \\</span>
<span class="sd">            m = \beta_1 * m + (1 - \beta_1) * g \\</span>
<span class="sd">            v = \beta_2 * v + (1 - \beta_2) * g * g \\</span>
<span class="sd">            update = \frac{m}{\sqrt{v} + \epsilon} \\</span>
<span class="sd">            update =</span>
<span class="sd">            \begin{cases}</span>
<span class="sd">                update + weight\_decay * w</span>
<span class="sd">                    &amp; \text{ if } weight\_decay &gt; 0 \\</span>
<span class="sd">                update</span>
<span class="sd">                    &amp; \text{ otherwise }</span>
<span class="sd">            \end{cases} \\</span>
<span class="sd">            w  = w - lr * update</span>
<span class="sd">        \end{array}</span>

<span class="sd">    :math:`m` represents the 1st moment vector, :math:`v` represents the 2nd moment vector, :math:`g` represents</span>
<span class="sd">    `gradient`, :math:`\beta_1, \beta_2` represent `beta1` and `beta2`,</span>
<span class="sd">    :math:`lr` represents `learning_rate`, :math:`w` represents `var`, :math:`decay` represents `weight_decay`,</span>
<span class="sd">    :math:`\epsilon` represents `epsilon`.</span>

<span class="sd">    Args:</span>
<span class="sd">        use_locking (bool): Whether to enable a lock to protect variable tensors from being updated.</span>
<span class="sd">            If true, updates of the var, m, and v tensors will be protected by a lock.</span>
<span class="sd">            If false, the result is unpredictable. Default: False.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **var** (Tensor) - Weights to be updated with the type float16 or float32.</span>
<span class="sd">        - **m** (Tensor) - The 1st moment vector in the updating formula with the type float32.</span>
<span class="sd">        - **v** (Tensor) - the 2nd moment vector in the updating formula with the type float32.</span>
<span class="sd">        - **lr** (float) - :math:`lr` in the updating formula.</span>
<span class="sd">        - **beta1** (float) - The exponential decay rate for the 1st moment estimations.</span>
<span class="sd">        - **beta2** (float) - The exponential decay rate for the 2nd moment estimations.</span>
<span class="sd">        - **epsilon** (float) - Term added to the denominator to improve numerical stability.</span>
<span class="sd">        - **decay** (float) - The weight decay value, must be a scalar tensor with float data type.</span>
<span class="sd">        - **gradient** (Tensor) - Gradient, has the type float16.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tuple of 3 Tensor, the updated parameters.</span>

<span class="sd">        - **var** (Tensor) - The same shape and data type as `var`.</span>
<span class="sd">        - **m** (Tensor) - The same shape and data type as `m`.</span>
<span class="sd">        - **v** (Tensor) - The same shape and data type as `v`.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.context as context</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.nn as nn</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, Parameter</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import dtype as mstype</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.opt = ops.FusedCastAdamWeightDecay()</span>
<span class="sd">        ...         self.var = Parameter(Tensor(np.ones([2, 2]), mstype.float16), name=&quot;var&quot;)</span>
<span class="sd">        ...         self.m = Parameter(Tensor(np.ones([2, 2]), mstype.float32), name=&quot;m&quot;)</span>
<span class="sd">        ...         self.v = Parameter(Tensor(np.ones([2, 2]), mstype.float32), name=&quot;v&quot;)</span>
<span class="sd">        ...     def construct(self, lr, beta1, beta2, epsilon, decay, grad):</span>
<span class="sd">        ...         out = self.opt(self.var, self.m, self.v, lr, beta1, beta2, epsilon, decay, grad)</span>
<span class="sd">        ...         return out</span>
<span class="sd">        &gt;&gt;&gt; context.set_context(mode=context.GRAPH_MODE, device_target=&quot;CPU&quot;)</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; gradient = Tensor(np.ones([2, 2]), mstype.float16)</span>
<span class="sd">        &gt;&gt;&gt; output = net(0.001, 0.9, 0.999, 1e-8, 0.0, gradient)</span>
<span class="sd">        &gt;&gt;&gt; print(net.var.asnumpy())</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">use_locking</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;side_effect_mem&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;use_locking&quot;</span><span class="p">,</span> <span class="n">use_locking</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">var_shape</span><span class="p">,</span> <span class="n">m_shape</span><span class="p">,</span> <span class="n">v_shape</span><span class="p">,</span> <span class="n">lr_shape</span><span class="p">,</span> <span class="n">beta1_shape</span><span class="p">,</span> <span class="n">beta2_shape</span><span class="p">,</span>
                    <span class="n">epsilon_shape</span><span class="p">,</span> <span class="n">decay_shape</span><span class="p">,</span> <span class="n">grad_shape</span><span class="p">,</span> <span class="n">global_norm</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;var_shape&quot;</span><span class="p">,</span> <span class="n">var_shape</span><span class="p">,</span> <span class="s2">&quot;m_shape&quot;</span><span class="p">,</span> <span class="n">m_shape</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;var_shape&quot;</span><span class="p">,</span> <span class="n">var_shape</span><span class="p">,</span> <span class="s2">&quot;v_shape&quot;</span><span class="p">,</span> <span class="n">v_shape</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;var_shape&quot;</span><span class="p">,</span> <span class="n">var_shape</span><span class="p">,</span> <span class="s2">&quot;grad_shape&quot;</span><span class="p">,</span> <span class="n">grad_shape</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">var_shape</span><span class="p">,</span> <span class="n">m_shape</span><span class="p">,</span> <span class="n">v_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">var_dtype</span><span class="p">,</span> <span class="n">m_dtype</span><span class="p">,</span> <span class="n">v_dtype</span><span class="p">,</span> <span class="n">lr_dtype</span><span class="p">,</span> <span class="n">beta1_dtype</span><span class="p">,</span> <span class="n">beta2_dtype</span><span class="p">,</span>
                    <span class="n">epsilon_dtype</span><span class="p">,</span> <span class="n">decay_dtype</span><span class="p">,</span> <span class="n">grad_dtype</span><span class="p">,</span> <span class="n">global_norm</span><span class="p">):</span>
        <span class="n">args</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;m&quot;</span><span class="p">:</span> <span class="n">m_dtype</span><span class="p">,</span> <span class="s2">&quot;v&quot;</span><span class="p">:</span> <span class="n">v_dtype</span><span class="p">}</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_tensors_dtypes_same_and_valid</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">number_type</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_scalar_or_tensor_types_same</span><span class="p">({</span><span class="s2">&quot;var&quot;</span><span class="p">:</span> <span class="n">var_dtype</span><span class="p">},</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_scalar_or_tensor_types_same</span><span class="p">({</span><span class="s2">&quot;grad&quot;</span><span class="p">:</span> <span class="n">grad_dtype</span><span class="p">},</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>

        <span class="n">args</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;lr&quot;</span><span class="p">:</span> <span class="n">lr_dtype</span><span class="p">,</span> <span class="s2">&quot;beta1&quot;</span><span class="p">:</span> <span class="n">beta1_dtype</span><span class="p">,</span> <span class="s2">&quot;beta2&quot;</span><span class="p">:</span> <span class="n">beta2_dtype</span><span class="p">,</span> <span class="s2">&quot;epsilon&quot;</span><span class="p">:</span> <span class="n">epsilon_dtype</span><span class="p">,</span>
                <span class="s2">&quot;decay&quot;</span><span class="p">:</span> <span class="n">decay_dtype</span><span class="p">}</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_scalar_or_tensor_types_same</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">var_dtype</span><span class="p">,</span> <span class="n">m_dtype</span><span class="p">,</span> <span class="n">v_dtype</span>


<span class="k">class</span> <span class="nc">FusedAdaFactor</span><span class="p">(</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Updates gradients by the Adaptive Learning Rates with Sublinear Memory Cost (Adafactor) algorithm.</span>

<span class="sd">    The Adafactor algorithm is proposed in `Adafactor: Adafactor: Adaptive Learning Rates with Sublinear Memory</span>
<span class="sd">    Cost &lt;https://arxiv.org/abs/1804.04235&gt;`_.</span>

<span class="sd">    .. warning::</span>
<span class="sd">        This is an experimental prototype that is subject to change and/or deletion.</span>

<span class="sd">    Adafactor for weight vector are as follows,</span>

<span class="sd">    .. math::</span>
<span class="sd">        \begin{array}{l} \\</span>
<span class="sd">        \alpha_{t}=\max \left(\epsilon_{2}, \operatorname{RMS}\left(X_{t-1}\right)\right) \rho_{t} \\</span>
<span class="sd">        G_{t}=\nabla f_{t}\left(X_{t-1}\right) \\</span>
<span class="sd">        \hat{V}_{t}=\hat{\beta}_{2} \hat{V}_{t-1}+\left(1-\hat{\beta}_{2_{t}}\right)\left(G_{t}^{2}+ \\</span>
<span class="sd">        \epsilon_{1} 1_{n}\right) \\</span>
<span class="sd">        U_{t}=G_{t} / \sqrt{\hat{V}_{t}} \\</span>
<span class="sd">        \hat{U}_{t}=U_{t} / \max \left(1, \operatorname{RMS}\left(U_{t}\right) / d\right) \\</span>
<span class="sd">        X_{t}=X_{t-1}-\alpha_{t} \hat{U}_{t}</span>
<span class="sd">        \end{array}</span>

<span class="sd">    Adafactor for weight matrices are as follows,</span>

<span class="sd">    .. math::</span>
<span class="sd">        \begin{array}{l} \\</span>
<span class="sd">        \alpha_{t}=\max \left(\epsilon_{2}, \operatorname{RMS}\left(X_{t-1}\right)\right) \rho_{t} \\</span>
<span class="sd">        G_{t}=\nabla f_{t}\left(X_{t-1}\right) \\</span>
<span class="sd">        R_{t}=\hat{\beta}_{2 t} R_{t-1}+\left(1-\hat{\beta}_{2 t}\right)\left(G_{t}^{2}+ \\</span>
<span class="sd">        \epsilon_{1} 1_{n} 1_{m}^{\top}\right) 1_{m} \\</span>
<span class="sd">        C_{t}=\hat{\beta}_{2 t} C_{t-1}+\left(1-\hat{\beta}_{2 t}\right) 1_{n}^{\top}\left(G_{t}^{2}+ \\</span>
<span class="sd">        \epsilon_{1} 1_{n} 1_{m}^{\top}\right) \\</span>
<span class="sd">        \hat{V}_{t}=R_{t} C_{t} / 1_{n}^{\top} R_{t} \\</span>
<span class="sd">        U_{t}=G_{t} / \sqrt{\hat{V}_{t}} \\</span>
<span class="sd">        \hat{U}_{t}=U_{t} / \max \left(1, \operatorname{RMS}\left(U_{t}\right) / d\right) \\</span>
<span class="sd">        X_{t}=X_{t-1}-\alpha_{t} U_{t}</span>
<span class="sd">        \end{array}</span>

<span class="sd">    Where RMS is:</span>

<span class="sd">    .. math::</span>
<span class="sd">        \operatorname{RMS}\left(U_{t}\right)=\operatorname{RMS}_{x \in X}\left(u_{x t}\right)= \\</span>
<span class="sd">        \sqrt{\operatorname{Mean}_{x \in X}\left(\frac{\left(g_{x t}\right)^{2}}{\hat{v}_{x t}}\right)}</span>

<span class="sd">    :math:`x` is each individual parameter,</span>
<span class="sd">    :math:`t` is assumed to be the current number of steps,</span>
<span class="sd">    :math:`a_{t}` is the learning rate,</span>
<span class="sd">    :math:`f(X)` is the loss function,</span>
<span class="sd">    :math:`\epsilon1` and :math:`\epsilon2` is a small positive number to prevent errors,</span>
<span class="sd">    :math:`d` is the clipping threshold,</span>
<span class="sd">    :math:`\beta_{2}` is the moment decay,</span>
<span class="sd">    :math:`\rho` is the relative step size,</span>
<span class="sd">    :math:`R` is the running averages of the row sums of the squared gradient,</span>
<span class="sd">    :math:`C` is the running averages of the column sums of the squared gradient.</span>

<span class="sd">    Args:</span>
<span class="sd">        enable_weight_decay (bool): If True, enable weight decay. default: False</span>
<span class="sd">        enable_first_moment (bool): If True, enable first moment. default: False</span>
<span class="sd">        enable_scale_parameter (bool): If True, enable scale learning rate using parameter. default: False</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **epsilon** (Tensor) - input epsilon pair.</span>
<span class="sd">        - **clip_threshold** (float) - The threshold of root mean square of final gradient update.</span>
<span class="sd">        - **beta1** (float) - The exponential decay rate for the 1nd moment estimations.</span>
<span class="sd">        - **beta2** (float) - The exponential decay rate for the 2nd moment estimations.</span>
<span class="sd">        - **weight_decay** (float) - The weight decay value, must be a scalar tensor with float data type.</span>
<span class="sd">        - **learning_rate** (float) - The learning rate value.</span>
<span class="sd">        - **gradient** (Tensor) - Gradient.</span>
<span class="sd">        - **param** (Tensor) - Weights to be updated.</span>
<span class="sd">        - **exp_avg** (Tensor) - The exponential moving average of 1st moment optimizer state.</span>
<span class="sd">        - **exp_avg_sq_row** (Tensor) - The exponential moving average of square of gradient square row factor.</span>
<span class="sd">        - **exp_avg_sq_col** (Tensor) - The exponential moving average of square of gradient square col factor.</span>
<span class="sd">        - **exp_avg_sq** (Tensor) - The exponential moving average of square of gradient square.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        - **dummy_param** (Tensor) - The same shape and data type as `param`.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; import numpy as np</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.context as context</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.nn as nn</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import Tensor, Parameter</span>
<span class="sd">        &gt;&gt;&gt; from mindspore import dtype as mstype</span>
<span class="sd">        &gt;&gt;&gt; param_shape = [2, 3, 2]</span>
<span class="sd">        &gt;&gt;&gt; class Net(nn.Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(Net, self).__init__()</span>
<span class="sd">        ...         self.opt = ops.FusedAdaFactor()</span>
<span class="sd">        ...         self.param = Parameter(Tensor(np.ones(param_shape), mstype.float32), name=&quot;param&quot;)</span>
<span class="sd">        ...         self.exp_avg = Parameter(Tensor(np.zeros(param_shape), mstype.float32), name=&quot;exp_avg&quot;)</span>
<span class="sd">        ...         self.exp_avg_sq = Parameter(Tensor(np.zeros(param_shape), mstype.float32), name=&quot;exp_avg_sq&quot;)</span>
<span class="sd">        ...         self.exp_avg_sq_row = Parameter(Tensor(np.zeros([2, 3]), mstype.float32), name=&quot;exp_avg_sq_row&quot;)</span>
<span class="sd">        ...         self.exp_avg_sq_col = Parameter(Tensor(np.zeros([2, 2]), mstype.float32), name=&quot;exp_avg_sq_col&quot;)</span>
<span class="sd">        ...</span>
<span class="sd">        ...     def construct(self, epsilon, clip_threshold, beta1, beta2, weight_decay, lr, grad):</span>
<span class="sd">        ...         out = self.opt(epsilon, clip_threshold, beta1, beta2, weight_decay, lr, grad, self.param,</span>
<span class="sd">        ...                        self.exp_avg, self.exp_avg_sq_row, self.exp_avg_sq_col, self.exp_avg_sq)</span>
<span class="sd">        ...         return out</span>
<span class="sd">        &gt;&gt;&gt; context.set_context(mode=context.GRAPH_MODE, device_target=&quot;CPU&quot;)</span>
<span class="sd">        &gt;&gt;&gt; net = Net()</span>
<span class="sd">        &gt;&gt;&gt; gradient = Tensor(np.ones(param_shape), mstype.float32)</span>
<span class="sd">        &gt;&gt;&gt; net((1e-30, 1e-3), 1.0, 0.9, 0.8, 1e-2, 0.03, gradient)</span>
<span class="sd">        &gt;&gt;&gt; print(net.param.asnumpy())</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@prim_attr_register</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">enable_scale_parameter</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">enable_first_moment</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">enable_weight_decay</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;side_effect_mem&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;enable_scale_parameter&quot;</span><span class="p">,</span> <span class="n">enable_scale_parameter</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;enable_first_moment&quot;</span><span class="p">,</span> <span class="n">enable_first_moment</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;enable_weight_decay&quot;</span><span class="p">,</span> <span class="n">enable_weight_decay</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">epsilon_shape</span><span class="p">,</span> <span class="n">clip_threshold_shape</span><span class="p">,</span> <span class="n">beta1_shape</span><span class="p">,</span> <span class="n">beta2t_shape</span><span class="p">,</span> <span class="n">weight_decay_shape</span><span class="p">,</span>
                    <span class="n">learning_rate_shape</span><span class="p">,</span> <span class="n">grad_shape</span><span class="p">,</span> <span class="n">param_shape</span><span class="p">,</span> <span class="n">exp_avg_shape</span><span class="p">,</span> <span class="n">exp_avg_sq_row_shape</span><span class="p">,</span>
                    <span class="n">exp_avg_sq_col_shape</span><span class="p">,</span> <span class="n">exp_avg_sq_shape</span><span class="p">):</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check</span><span class="p">(</span><span class="s2">&quot;grad_shape&quot;</span><span class="p">,</span> <span class="n">grad_shape</span><span class="p">,</span> <span class="s2">&quot;param_shape&quot;</span><span class="p">,</span> <span class="n">param_shape</span><span class="p">,</span> <span class="n">Rel</span><span class="o">.</span><span class="n">EQ</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">param_shape</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">epsilon_type</span><span class="p">,</span> <span class="n">clip_threshold_type</span><span class="p">,</span> <span class="n">beta1_type</span><span class="p">,</span> <span class="n">beta2t_type</span><span class="p">,</span> <span class="n">weight_decay_type</span><span class="p">,</span>
                    <span class="n">learning_rate_type</span><span class="p">,</span> <span class="n">grad_type</span><span class="p">,</span> <span class="n">param_type</span><span class="p">,</span> <span class="n">exp_avg_type</span><span class="p">,</span> <span class="n">exp_avg_sq_row_type</span><span class="p">,</span>
                    <span class="n">exp_avg_sq_col_type</span><span class="p">,</span> <span class="n">exp_avg_sq_type</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">param_type</span>
</pre></div>

           </div>
           
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>
        &#169; Copyright 2021, MindSpore.

    </p>
  </div>
    
    
    
    Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    
    provided by <a href="https://readthedocs.org">Read the Docs</a>. 

</footer>
        </div>
      </div>

    </section>

  </div>
  

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
   

</body>
</html>